import tkinter as tk
from tkinter import ttk, scrolledtext, filedialog, messagebox, simpledialog
import threading
import queue
import os
import json
from datetime import datetime
# pyinstaller DeepSeekChatGUI.spec 
# 导入你的 DeepSeekChat 类
from DeepSeekChat import DeepSeekChat

class DeepSeekChatGUI:
    def __init__(self, root):
        self.root = root
        self.root.title("DeepSeek 小说创作助手")
        self.root.geometry("1000x700")
        self.root.configure(bg='#f0f0f0')
        self.is_ai_response_started = False
        
        # 配置文件路径
        self.config_file = "config.json"
        
        # 尝试加载保存的API Key
        self.api_key = self.load_api_key()
        
        # 如果未找到保存的API Key，则提示用户输入
        if self.api_key is None or len(self.api_key) == 0:
            self.api_key = self.prompt_api_key()
            if self.api_key is None:
                self.root.destroy()
                return
            # 询问用户是否要保存API Key
            if messagebox.askyesno("保存API Key", "是否要保存API Key以便下次使用？"):
                self.save_api_key(self.api_key)
            
        # 初始化 DeepSeekChat
        self.chat_client = DeepSeekChat(api_key=self.api_key)
        
        # 创建界面
        self.create_widgets()
        self.setup_layout()
        
        # 流式响应队列
        self.stream_queue = queue.Queue()
        self.is_streaming = False
        
        # 启动队列检查
        self.check_queue()
        
        # 刷新系统提示词列表
        self.refresh_system_prompts()
    
    def load_api_key(self):
        """从配置文件加载API Key"""
        try:
            if os.path.exists(self.config_file):
                with open(self.config_file, 'r') as f:
                    config = json.load(f)
                    return config.get('api_key')
        except Exception as e:
            print(f"加载配置文件失败: {str(e)}")
        return None
    
    def save_api_key(self, api_key):
        """保存API Key到配置文件"""
        try:
            config = {'api_key': api_key}
            with open(self.config_file, 'w') as f:
                json.dump(config, f)
            return True
        except Exception as e:
            messagebox.showerror("错误", f"保存API Key失败: {str(e)}")
            return False
    
    def clear_saved_api_key(self):
        """清除已保存的API Key"""
        try:
            if os.path.exists(self.config_file):
                os.remove(self.config_file)
                return True
        except Exception as e:
            messagebox.showerror("错误", f"清除API Key失败: {str(e)}")
        return False
    
    def prompt_api_key(self):
        """弹出API Key输入对话框"""
        api_key = simpledialog.askstring(
            "API Key 输入", 
            "请输入DeepSeek API Key:", 
            show='*', 
            parent=self.root
        )
        
        if api_key is None or api_key.strip() == '':
            messagebox.showerror("错误", "必须提供有效的API Key才能使用本程序")
            return None
            
        return api_key.strip()
    
    def update_api_key(self):
        """更新API Key"""
        new_api_key = self.prompt_api_key()
        if new_api_key is None:
            return
            
        # 确认是否要重新初始化
        if messagebox.askyesno("确认", "更改API Key将重置当前对话，是否继续？"):
            self.api_key = new_api_key
            
            # 询问是否保存新的API Key
            if messagebox.askyesno("保存API Key", "是否要保存新的API Key以便下次使用？"):
                if self.save_api_key(new_api_key):
                    self.status_var.set("API Key已更新并保存")
                else:
                    self.status_var.set("API Key已更新但保存失败")
            else:
                # 如果不保存，清除之前保存的
                self.clear_saved_api_key()
                self.status_var.set("API Key已更新（未保存）")
            
            # 重新初始化DeepSeekChat
            self.chat_client = DeepSeekChat(api_key=self.api_key)
            # 清空聊天历史
            self.chat_history.config(state=tk.NORMAL)
            self.chat_history.delete("1.0", tk.END)
            self.chat_history.config(state=tk.DISABLED)
            # 刷新系统提示词
            self.refresh_system_prompts()
            # 刷新参考资料列表
            self.refresh_reference_list()
    
    def create_widgets(self):
        # 主框架
        self.main_frame = ttk.Frame(self.root, padding=10)
        
        # 标题
        self.title_label = ttk.Label(
            self.main_frame, 
            text="DeepSeek 小说创作助手", 
            font=("Arial", 16, "bold")
        )
        
        # 系统提示词区域
        self.system_frame = ttk.LabelFrame(self.main_frame, text="系统提示词管理", padding=5)
        
        # 系统提示词选择
        self.system_label = ttk.Label(self.system_frame, text="当前系统提示词:")
        self.system_var = tk.StringVar()
        self.system_combo = ttk.Combobox(self.system_frame, textvariable=self.system_var, state="readonly")
        self.system_combo.bind('<<ComboboxSelected>>', self.on_system_selected)
        
        self.system_buttons_frame = ttk.Frame(self.system_frame)
        self.system_create_button = ttk.Button(self.system_buttons_frame, text="新建", command=self.create_system_prompt)
        self.system_edit_button = ttk.Button(self.system_buttons_frame, text="编辑", command=self.edit_system_prompt)
        self.system_delete_button = ttk.Button(self.system_buttons_frame, text="删除", command=self.delete_system_prompt)
        self.system_refresh_button = ttk.Button(self.system_buttons_frame, text="刷新", command=self.refresh_system_prompts)
        
        # 参考资料区域
        self.ref_frame = ttk.LabelFrame(self.main_frame, text="参考资料管理", padding=5)
        self.ref_listbox = tk.Listbox(self.ref_frame, height=8, width=40)
        self.ref_scrollbar = ttk.Scrollbar(self.ref_frame, orient=tk.VERTICAL, command=self.ref_listbox.yview)
        self.ref_listbox.configure(yscrollcommand=self.ref_scrollbar.set)

        self.ref_buttons_frame = ttk.Frame(self.ref_frame)
        # 修改按钮布局从水平变为垂直
        self.add_ref_button = ttk.Button(self.ref_buttons_frame, text="添加文件", command=self.add_reference_file)
        self.view_ref_button = ttk.Button(self.ref_buttons_frame, text="查看内容", command=self.view_reference_file)
        self.del_ref_button = ttk.Button(self.ref_buttons_frame, text="删除文件", command=self.delete_reference_file)
        self.refresh_ref_button = ttk.Button(self.ref_buttons_frame, text="刷新列表", command=self.refresh_reference_list)
        
        # 对话历史区域
        self.chat_frame = ttk.LabelFrame(self.main_frame, text="对话历史", padding=5)
        self.chat_history = scrolledtext.ScrolledText(
            self.chat_frame,
            wrap=tk.WORD,
            width=80,
            height=25,
            font=("Arial", 11),
            state=tk.DISABLED
        )
        
        # 输入区域
        self.input_frame = ttk.Frame(self.main_frame)
        self.user_input = scrolledtext.ScrolledText(
            self.input_frame,
            wrap=tk.WORD,
            width=70,
            height=4,
            font=("Arial", 11)
        )
        
        # 控制按钮区域
        self.control_frame = ttk.Frame(self.main_frame)
        
        self.use_ref_var = tk.BooleanVar(value=True)
        self.use_ref_check = ttk.Checkbutton(
            self.control_frame, 
            text="使用参考资料", 
            variable=self.use_ref_var
        )
        
        self.send_button = ttk.Button(
            self.control_frame,
            text="发送",
            command=self.send_message
        )
        
        self.stop_button = ttk.Button(
            self.control_frame,
            text="停止生成",
            command=self.stop_generation,
            state=tk.DISABLED
        )
        
        # 功能按钮区域
        self.tools_frame = ttk.Frame(self.main_frame)
        
        self.save_button = ttk.Button(
            self.tools_frame,
            text="保存对话",
            command=self.save_conversation
        )
        
        self.export_button = ttk.Button(
            self.tools_frame,
            text="导出对话",
            command=self.export_conversation
        )
        
        self.clear_button = ttk.Button(
            self.tools_frame,
            text="清空对话",
            command=self.clear_conversation
        )
        
        self.new_chat_button = ttk.Button(
            self.tools_frame,
            text="新建对话",
            command=self.new_conversation
        )
        
        # API Key 按钮
        self.api_key_button = ttk.Button(
            self.tools_frame,
            text="设置API Key",
            command=self.update_api_key
        )
        
        # 清除API Key按钮
        self.clear_api_key_button = ttk.Button(
            self.tools_frame,
            text="清除API Key",
            command=self.clear_api_key
        )
        
        # 状态栏
        self.status_var = tk.StringVar(value="就绪")
        self.status_bar = ttk.Label(
            self.root, 
            textvariable=self.status_var, 
            relief=tk.SUNKEN, 
            anchor=tk.W
        )
    
    def clear_api_key(self):
        """清除已保存的API Key"""
        if messagebox.askyesno("确认", "确定要清除已保存的API Key吗？下次启动时需要重新输入。"):
            if self.clear_saved_api_key():
                self.status_var.set("已清除保存的API Key")
            else:
                self.status_var.set("清除API Key失败")
    
    def setup_layout(self):
        # 主框架布局
        self.main_frame.pack(fill=tk.BOTH, expand=True)
        
        # 标题
        self.title_label.grid(row=0, column=0, columnspan=2, pady=(0, 15))
        
        # 系统提示词区域
        self.system_frame.grid(row=1, column=0, columnspan=2, pady=(0, 10), sticky=tk.EW)
        self.system_label.grid(row=0, column=0, padx=(0, 5), sticky=tk.W)
        self.system_combo.grid(row=0, column=1, padx=(0, 10), sticky=tk.EW)
        self.system_buttons_frame.grid(row=0, column=2, sticky=tk.EW)
        self.system_create_button.pack(side=tk.LEFT, padx=2)
        self.system_edit_button.pack(side=tk.LEFT, padx=2)
        self.system_delete_button.pack(side=tk.LEFT, padx=2)
        self.system_refresh_button.pack(side=tk.LEFT, padx=2)
        
        # 参考资料区域
        self.ref_frame.grid(row=2, column=0, rowspan=2, padx=(0, 10), sticky=tk.NSEW)
        self.ref_listbox.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)
        self.ref_scrollbar.pack(side=tk.RIGHT, fill=tk.Y)
        self.ref_buttons_frame.pack(fill=tk.BOTH, pady=(5, 0))
        self.add_ref_button.pack(side=tk.TOP, fill=tk.X, padx=2, pady=2)
        self.view_ref_button.pack(side=tk.TOP, fill=tk.X, padx=2, pady=2)
        self.del_ref_button.pack(side=tk.TOP, fill=tk.X, padx=2, pady=2)
        self.refresh_ref_button.pack(side=tk.TOP, fill=tk.X, padx=2, pady=2)
        
        # 对话历史区域
        self.chat_frame.grid(row=2, column=1, sticky=tk.NSEW)
        self.chat_history.pack(fill=tk.BOTH, expand=True)
        
        # 输入区域
        self.input_frame.grid(row=3, column=1, pady=(10, 5), sticky=tk.EW)
        self.user_input.pack(fill=tk.BOTH, expand=True)
        
        # 控制按钮区域
        self.control_frame.grid(row=4, column=0, columnspan=2, pady=(5, 10), sticky=tk.W)
        self.use_ref_check.pack(side=tk.LEFT, padx=(0, 10))
        self.send_button.pack(side=tk.LEFT, padx=(0, 10))
        self.stop_button.pack(side=tk.LEFT)
        
        # 功能按钮区域
        self.tools_frame.grid(row=5, column=0, columnspan=2, sticky=tk.W)
        self.save_button.pack(side=tk.LEFT, padx=(0, 10))
        self.export_button.pack(side=tk.LEFT, padx=(0, 10))
        self.clear_button.pack(side=tk.LEFT, padx=(0, 10))
        self.new_chat_button.pack(side=tk.LEFT, padx=(0, 10))
        self.api_key_button.pack(side=tk.LEFT, padx=(0, 10))
        self.clear_api_key_button.pack(side=tk.LEFT)
        
        # 状态栏
        self.status_bar.pack(side=tk.BOTTOM, fill=tk.X)
        
        # 配置网格权重
        self.main_frame.grid_rowconfigure(2, weight=1)
        self.main_frame.grid_columnconfigure(1, weight=1)
        self.system_frame.grid_columnconfigure(1, weight=1)
        
        # 初始化参考资料列表
        self.refresh_reference_list()
        
        # 绑定快捷键
        self.root.bind("<Control-Return>", lambda event: self.send_message())
    
    def refresh_system_prompts(self):
        """刷新系统提示词列表"""
        prompts = self.chat_client.list_system_prompts()
        self.system_combo['values'] = prompts
        
        # 设置当前选中的系统提示词
        current_system = self.chat_client.get_current_system_prompt_name()
        if current_system:
            self.system_var.set(current_system)
        elif prompts:
            self.system_var.set(prompts[0])
            self.chat_client.set_current_system_prompt(prompts[0])
    
    def on_system_selected(self, event):
        """当选择系统提示词时"""
        selected_system = self.system_var.get()
        if selected_system:
            self.chat_client.set_current_system_prompt(selected_system)
            # 清空聊天历史显示
            self.chat_history.config(state=tk.NORMAL)
            self.chat_history.delete("1.0", tk.END)
            self.chat_history.config(state=tk.DISABLED)
            self.status_var.set(f"已切换到系统提示词: {selected_system}")
    
    def create_system_prompt(self):
        """创建新系统提示词"""
        name = simpledialog.askstring("新建系统提示词", "请输入系统提示词名称:")
        if name:
            # 创建编辑窗口
            self.edit_system_prompt_window(name, "")
    
    def edit_system_prompt(self):
        """编辑选中的系统提示词"""
        selected_system = self.system_var.get()
        if not selected_system:
            messagebox.showwarning("警告", "请先选择一个系统提示词")
            return
            
        content = self.chat_client.load_system_prompt(selected_system)
        if content is not None:
            self.edit_system_prompt_window(selected_system, content)
        else:
            messagebox.showerror("错误", f"无法加载系统提示词: {selected_system}")
    
    def edit_system_prompt_window(self, name, content):
        """系统提示词编辑窗口"""
        # 创建编辑窗口
        edit_window = tk.Toplevel(self.root)
        edit_window.title(f"编辑系统提示词: {name}")
        edit_window.geometry("800x500")
        
        # 创建文本编辑区域
        prompt_label = ttk.Label(edit_window, text="系统提示词内容:")
        prompt_label.pack(pady=(10, 5), padx=10, anchor=tk.W)
        
        prompt_text = scrolledtext.ScrolledText(
            edit_window,
            wrap=tk.WORD,
            font=("Arial", 11),
            height=15
        )
        prompt_text.pack(fill=tk.BOTH, expand=True, padx=10, pady=5)
        prompt_text.insert(tk.END, content)
        
        # 按钮框架
        button_frame = ttk.Frame(edit_window)
        button_frame.pack(pady=10)
        
        def save_prompt():
            new_content = prompt_text.get("1.0", tk.END).strip()
            if self.chat_client.save_system_prompt(name, new_content):
                self.refresh_system_prompts()
                # 如果编辑的是当前选中的提示词，更新当前内容
                if self.chat_client.get_current_system_prompt_name() == name:
                    self.chat_client.current_system_prompt = new_content
                messagebox.showinfo("成功", "系统提示词已保存")
                edit_window.destroy()
        
        cancel_button = ttk.Button(button_frame, text="取消", command=edit_window.destroy)
        cancel_button.pack(side=tk.LEFT, padx=5)
        
        save_button = ttk.Button(button_frame, text="保存", command=save_prompt)
        save_button.pack(side=tk.LEFT, padx=5)
    
    def delete_system_prompt(self):
        """删除选中的系统提示词"""
        selected_system = self.system_var.get()
        if not selected_system:
            messagebox.showwarning("警告", "请先选择一个系统提示词")
            return
            
        if messagebox.askyesno("确认", f"确定要删除系统提示词 '{selected_system}' 吗？"):
            if self.chat_client.delete_system_prompt(selected_system):
                self.refresh_system_prompts()
                messagebox.showinfo("成功", "系统提示词已删除")
    
    def refresh_reference_list(self):
        """刷新参考资料列表"""
        self.ref_listbox.delete(0, tk.END)
        ref_files = self.chat_client.list_reference_files()
        for filename in ref_files:
            self.ref_listbox.insert(tk.END, filename)
    
    def add_reference_file(self):
        """添加参考资料文件"""
        file_path = filedialog.askopenfilename(
            title="选择参考资料文件",
            filetypes=[
                ("Text files", "*.txt"), 
                ("Python files", "*.py"),
                ("Markdown files", "*.md"),
                ("JSON files", "*.json"),
                ("All files", "*.*")
            ]
        )
        
        if file_path:
            # 将文件复制到参考资料目录
            import shutil
            filename = os.path.basename(file_path)
            dest_path = os.path.join(self.chat_client.reference_dir, filename)
            
            try:
                shutil.copy2(file_path, dest_path)
                self.refresh_reference_list()
                self.status_var.set(f"已添加参考资料: {filename}")
            except Exception as e:
                messagebox.showerror("错误", f"添加文件失败: {str(e)}")
    
    def view_reference_file(self):
        """查看选中的参考资料内容"""
        selection = self.ref_listbox.curselection()
        if not selection:
            messagebox.showwarning("警告", "请先选择一个参考资料文件")
            return
            
        filename = self.ref_listbox.get(selection[0])
        filepath = os.path.join(self.chat_client.reference_dir, filename)
        
        try:
            with open(filepath, "r", encoding="utf-8") as f:
                content = f.read()
                
            # 创建新窗口显示内容
            view_window = tk.Toplevel(self.root)
            view_window.title(f"参考资料: {filename}")
            view_window.geometry("600x400")
            
            text_widget = scrolledtext.ScrolledText(
                view_window,
                wrap=tk.WORD,
                font=("Arial", 11)
            )
            text_widget.pack(fill=tk.BOTH, expand=True, padx=10, pady=10)
            text_widget.insert(tk.END, content)
            text_widget.config(state=tk.DISABLED)
            
        except Exception as e:
            messagebox.showerror("错误", f"读取文件失败: {str(e)}")
    
    def delete_reference_file(self):
        """删除选中的参考资料文件"""
        selection = self.ref_listbox.curselection()
        if not selection:
            messagebox.showwarning("警告", "请先选择一个参考资料文件")
            return
            
        filename = self.ref_listbox.get(selection[0])
        filepath = os.path.join(self.chat_client.reference_dir, filename)
        
        if messagebox.askyesno("确认", f"确定要删除 {filename} 吗？"):
            try:
                os.remove(filepath)
                self.refresh_reference_list()
                self.status_var.set(f"已删除参考资料: {filename}")
            except Exception as e:
                messagebox.showerror("错误", f"删除文件失败: {str(e)}")
    
    def send_message(self):
        """发送消息给AI"""
        message = self.user_input.get("1.0", tk.END).strip()
        if not message:
            return
        
        # 清空输入框
        self.user_input.delete("1.0", tk.END)
    
        # 显示用户消息
        self.display_message("用户", message)
    
        # 重置AI响应开始标志
        self.is_ai_response_started = False
    
        # 更新按钮状态
        self.send_button.config(state=tk.DISABLED)
        self.stop_button.config(state=tk.NORMAL)
        self.is_streaming = True
    
        # 在新线程中获取AI回复
        use_ref = self.use_ref_var.get()
        threading.Thread(
            target=self.get_ai_response, 
            args=(message, use_ref), 
            daemon=True
        ).start()
    
    def get_ai_response(self, message, use_ref):
        """获取AI回复（在线程中运行）"""
        try:
            # 使用Stream API获取回复
            for chunk in self.chat_client.chat_stream(message, use_references=use_ref):
                self.stream_queue.put(chunk)
            
            # 流结束标志
            self.stream_queue.put(None)
            
        except Exception as e:
            error_msg = f"获取AI回复时出错: {str(e)}"
            self.stream_queue.put(error_msg)
            self.stream_queue.put(None)
    
    def check_queue(self):
        """检查队列并更新界面"""
        try:
            while True:
                chunk = self.stream_queue.get_nowait()
                if chunk is None:
                    # 流结束
                    self.is_streaming = False
                    self.send_button.config(state=tk.NORMAL)
                    self.stop_button.config(state=tk.DISABLED)
                    self.status_var.set("回复完成")
                
                    # 在AI回复结束后添加两个换行
                    self.chat_history.config(state=tk.NORMAL)
                    self.chat_history.insert(tk.END, "\n\n")
                    self.chat_history.see(tk.END)
                    self.chat_history.config(state=tk.DISABLED)
                
                    break
                else:
                    # 更新聊天历史
                    self.append_to_chat_history(chunk)
        except queue.Empty:
            pass
    
        # 继续检查队列
        self.root.after(100, self.check_queue)
    
    def display_message(self, sender, message):
        """显示消息到聊天历史"""
        self.chat_history.config(state=tk.NORMAL)
        self.chat_history.insert(tk.END, f"{sender}: {message}\n\n")
        self.chat_history.see(tk.END)
        self.chat_history.config(state=tk.DISABLED)
    
    def append_to_chat_history(self, chunk):
        """追加流式回复到聊天历史"""
        self.chat_history.config(state=tk.NORMAL)
    
        # 如果是AI响应的第一个块，确保前面有换行
        if not hasattr(self, 'is_ai_response_started'):
            self.is_ai_response_started = True
            # 检查是否需要添加换行
            current_content = self.chat_history.get("1.0", tk.END).strip()
            if current_content and not current_content.endswith('\n\n'):
                self.chat_history.insert(tk.END, "\n\n")
    
        self.chat_history.insert(tk.END, chunk)
        self.chat_history.see(tk.END)
        self.chat_history.config(state=tk.DISABLED)
    
    def stop_generation(self):
        """停止当前生成"""
        if self.is_streaming:
            self.chat_client.stop_current_stream()
            self.is_streaming = False
            self.send_button.config(state=tk.NORMAL)
            self.stop_button.config(state=tk.DISABLED)
            self.status_var.set("已停止生成")
    
    def save_conversation(self):
        """保存当前对话"""
        try:
            # 获取当前聊天内容
            self.chat_history.config(state=tk.NORMAL)
            content = self.chat_history.get("1.0", tk.END)
            self.chat_history.config(state=tk.DISABLED)
            
            # 弹出保存对话框
            file_path = filedialog.asksaveasfilename(
                defaultextension=".txt",
                filetypes=[("Text files", "*.txt"), ("All files", "*.*")]
            )
            
            if file_path:
                with open(file_path, "w", encoding="utf-8") as f:
                    f.write(content)
                self.status_var.set(f"对话已保存到: {file_path}")
                
        except Exception as e:
            messagebox.showerror("错误", f"保存失败: {str(e)}")
    
    def export_conversation(self):
        """导出完整对话记录"""
        try:
            export_file = self.chat_client.export_conversation()
            self.status_var.set(f"对话已导出到: {export_file}")
            messagebox.showinfo("成功", f"对话已导出到: {export_file}")
        except Exception as e:
            messagebox.showerror("错误", f"导出失败: {str(e)}")
    
    def clear_conversation(self):
        """清空当前对话"""
        if messagebox.askyesno("确认", "确定要清空当前对话吗？"):
            self.chat_client.clear_history()
            self.chat_history.config(state=tk.NORMAL)
            self.chat_history.delete("1.0", tk.END)
            self.chat_history.config(state=tk.DISABLED)
            self.status_var.set("对话已清空")
    
    def new_conversation(self):
        """创建新的对话文件"""
        if messagebox.askyesno("确认", "确定要创建新的对话文件吗？当前对话将继续显示但会保存到新文件"):
            self.chat_client.create_conversation_file()
            self.status_var.set("已创建新的对话文件")

def main():
    root = tk.Tk()
    app = DeepSeekChatGUI(root)
    root.mainloop()

# 创建主窗口并运行应用
if __name__ == "__main__":
    main()